{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Short Cut Solver"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this notebook make a basic demonstration of the ShortCutSolver. Use simplest set-up: Cartesian grid, squared Euclidean distance for cost, dense data structures (algorithm will only solve sparse sub-problems though)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# load libraries\n",
    "from lib.header_notebook import *\n",
    "\n",
    "# assume that CPLEX back-end is installed\n",
    "import Solvers.ShortCutSolver as ShortCutSolver\n",
    "import Solvers.OT_CPLEX as OT_CPLEX\n",
    "import Solvers.ShortCutSolver_CPLEX as ShortCutSolver_CPLEX\n",
    "\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Problem Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# load two 64x64 images\n",
    "imgX=sciio.loadmat(\"data/density-img/f-000-064.mat\")[\"a\"]\n",
    "imgY=sciio.loadmat(\"data/density-img/f-001-064.mat\")[\"a\"]\n",
    "\n",
    "# preprocessing (add small constant background mass, normalize masses, extract geometric positions of pixels)\n",
    "(muX,posX)=OTTools.processDensity_Grid(imgX,totalMass=1.,constOffset=1E-7)\n",
    "(muY,posY)=OTTools.processDensity_Grid(imgY,totalMass=1.,constOffset=1E-7)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# visualize images\n",
    "fig=plt.figure(figsize=(8,4))\n",
    "fig.add_subplot(1,2,1)\n",
    "plt.imshow(imgX)\n",
    "fig.add_subplot(1,2,2)\n",
    "plt.imshow(imgY)\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# set up hierarchical partitions\n",
    "\n",
    "# finest layer above image has 2^partitionDepth grid points per dimension, then one below is image\n",
    "partitionDepth=5\n",
    "# another partition parameter, to be discussed later\n",
    "partitionChildMode=HierarchicalPartition.THPMode_Grid\n",
    "\n",
    "# create partitions from point clouds & measures, export partitions already to c++ library for later use\n",
    "(partitionX,pointerX)=HierarchicalPartition.GetPartition(posX,partitionDepth,partitionChildMode,imgX.shape, mu=muX,\\\n",
    "    signal_pos=True, signal_radii=False,clib=SolverCFC, export=True, verbose=False)\n",
    "\n",
    "(partitionY,pointerY)=HierarchicalPartition.GetPartition(posY,partitionDepth,partitionChildMode,imgY.shape, mu=muY,\\\n",
    "    signal_pos=True, signal_radii=True,clib=SolverCFC, export=True, verbose=False)\n",
    "\n",
    "pointerYpos=HierarchicalPartition.getSignalPointer(partitionY,\"pos\")\n",
    "pointerYradii=HierarchicalPartition.getSignalPointer(partitionY,\"radii\", lBottom=partitionY.nlayers-2)\n",
    "\n",
    "\n",
    "# for demonstration purposes: compute dense cost functions at all levels\n",
    "p=2. # exponent for Euclidean distance\n",
    "cList=HierarchicalPartition.GetHierarchicalCost(partitionX, partitionY,\\\n",
    "    lambda posx, posy : OTTools.getEuclideanCostFunction(posx,posy,p=p))\n",
    "\n",
    "# print a few stats on the created problem\n",
    "print(\"cells in partition x: \", partitionX.cardLayers)\n",
    "print(\"cells in partition y: \", partitionY.cardLayers)\n",
    "print(\"dimensions of dense costs: \",[c.shape for c in cList])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Solving"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Multiscale Solving"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# the algorithm has a modular structure. different components can be combined to final algorithm\n",
    "\n",
    "# refinement:\n",
    "#     generate initial fine coupling support when doing a layer refinement\n",
    "methodSetup_Refinement=ShortCutSolver.getMethodSetup_Refinement(pointerX,pointerY,SolverCFC)\n",
    "\n",
    "# coupling handler:\n",
    "#     data structure for handling cost function and coupling.\n",
    "#     in this small example use naive dense data structures (so no memory will be saved, only run-time)\n",
    "#     for a sparse data structure see other example files.\n",
    "\n",
    "methodSetup_CouplingHandler=ShortCutSolver.getMethodSetup_CouplingHandler_SemiDense()\n",
    "\n",
    "# solver for sparse sub-problems\n",
    "#     in this example only use CPLEX solver (Lemon requires some more preprocessing of densities)\n",
    "#     couplingHandlerType must match the coupling handler chosen above (nothing to worry about now)\n",
    "#     initializeBases=True indicates that warm-starting the solver during iterations on same scale will be used.\n",
    "methodSetup_SubSolver=ShortCutSolver_CPLEX.getMethodSetup_SubSolver_CPLEX(\\\n",
    "        couplingHandlerType=ShortCutSolver_CPLEX.CH_SemiDense,initializeBases=True)\n",
    "\n",
    "\n",
    "# shielding\n",
    "#     one must choose a shielding method that matches the geometry of the cost\n",
    "#     in this example we have the simplest possible case: squared Euclidean distance on regular Cartesian grid.\n",
    "#     so can choose simplest (and fastest) shielding method\n",
    "\n",
    "methodSetup_Shielding=ShortCutSolver.getMethodSetup_Shielding_Grid()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# do multi-scale solving.\n",
    "#     algorithm gets all the chosen methods above and combines them into full ShortCut solver.\n",
    "#     solves successively from very coarse level to finest level\n",
    "#     is configured in verbose mode. at each level a small report is printed\n",
    "time1=datetime.datetime.now()\n",
    "result=ShortCutSolver.MultiscaleSolver(partitionX, partitionY, cList,\\\n",
    "    methodSetup_Refinement, methodSetup_CouplingHandler, methodSetup_SubSolver, methodSetup_Shielding,\\\n",
    "    nLayerInitial=1,nLayerFinal=None,\\\n",
    "    Verbose=True,\\\n",
    "    maxSteps=100,collectReports=True,measureTimes=True,stepwiseAnalysis=False)\n",
    "time2=datetime.datetime.now()\n",
    "print(time2-time1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Verify Shielding Condition and Optimality"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# extract final neighbourhood and support from solver\n",
    "xVars=ShortCutSolver.SolverGetXVars(result[0][\"pointer\"])\n",
    "xSupport=ShortCutSolver.SolverGetSupport(result[0][\"pointer\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# verify shielding condition\n",
    "\n",
    "## do some preprocessing to gather required information\n",
    "\n",
    "# extract some assignment map\n",
    "xMap=xSupport[0][xSupport[1][:-1]]\n",
    "\n",
    "# generate neighbourhoods\n",
    "neighboursList=[None]\n",
    "for i in range(1,partitionX.nlayers):\n",
    "    neighboursList.append(ShortCutSolver.GetGridNeighbours(partitionX.dims[i]))\n",
    "\n",
    "nLayer=partitionX.nlayers-1\n",
    "shielding_info=ShortCutSolver.VerifyShielding(cList[nLayer],xVars[0],xVars[1],\\\n",
    "        neighboursList[nLayer][0],neighboursList[nLayer][1],xMap)\n",
    "\n",
    "print(\"shielding condition test:\", shielding_info[0])\n",
    "print(\"missed shields:\", shielding_info[1][0].shape[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# test explicitly for violated constraints\n",
    "constraint_info=ShortCutSolver.VerifyDualConstraints(cList[nLayer],\\\n",
    "    result[0][\"result_subsolver\"][\"alpha\"],result[0][\"result_subsolver\"][\"beta\"],1E-10)\n",
    "\n",
    "print(\"constraint violation test:\", constraint_info[0])\n",
    "print(\"violated constraints:\", constraint_info[1][0].shape[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Check Objective"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "mu=result[0][\"result_couplinghandler\"][\"mu\"]\n",
    "alpha=result[0][\"result_subsolver\"][\"alpha\"]\n",
    "beta=result[0][\"result_subsolver\"][\"beta\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# primal score\n",
    "np.sum(mu*cList[-1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# dual score\n",
    "np.sum(alpha*muX)+np.sum(beta*muY)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Dense Solver"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "nLayer=partitionX.nlayers-1\n",
    "muX=partitionX.layers[nLayer][\"mass\"]\n",
    "muY=partitionY.layers[nLayer][\"mass\"]\n",
    "\n",
    "time1=datetime.datetime.now()\n",
    "resultDense_CPLEX=OT_CPLEX.solveDense(cList[nLayer],muX,muY)\n",
    "time2=datetime.datetime.now()\n",
    "deltaTDense=(time2-time1).total_seconds()\n",
    "print(time2-time1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# dense primal score\n",
    "np.sum(resultDense_CPLEX[\"mu\"]*cList[nLayer])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Cleaning Up"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "ShortCutSolver.SolverClose(result[0][\"pointer\"])\n",
    "SolverCFC.Close(pointerX)\n",
    "SolverCFC.Close(pointerY)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.4.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
